import tqdm
from itertools import chain
import numpy as np
from scipy.spatial.transform import Rotation
from pose import initial_pose
import argparse
import cv2
from scipy.optimize import least_squares


def arg_parse():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument("-f", type=float)
    arg_parser.add_argument("-d", type=float, default=5.5)
    arg_parser.add_argument("--reso-x", type=int, default=4504)
    arg_parser.add_argument("--reso-y", type=int, default=4505)
    arg_parser.add_argument("--noise", type=float, default=0)
    arg_parser.add_argument("--drift", type=float, default=0)
    arg_parser.add_argument("-m", type=int, default=50)
    arg_parser.add_argument("--height", type=int, default=300)
    arg_parser.add_argument("-n", type=int, default=100)
    arg_parser.add_argument("--outlier", type=float, default=0)
    return arg_parser.parse_args()


def linear_least(real_points, observations, rotation_matrix, K):
    Y = (
        np.linalg.inv(K @ np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1]]))
        @ np.pad(observations, ((0, 0), (0, 1)), "constant", constant_values=1).T
    )  # 2 x N
    Yx = np.stack([np.zeros_like(Y[0]), -Y[2], Y[1]])  # 3 x N
    Yy = np.stack([Y[2], np.zeros_like(Y[0]), -Y[0]])
    Yz = np.stack([-Y[1], Y[0], np.zeros_like(Y[0])])
    Y_outer = np.concatenate([Yx, Yy, Yz], axis=1)  # 3 x 2N
    X_outer = np.concatenate([rotation_matrix @ real_points.T] * 3, axis=1)  # 3 x 2N
    translation = np.linalg.lstsq(
        Y_outer.T, -np.sum(Y_outer * X_outer, axis=0), rcond=1
    )[0]
    return translation


def non_linear_least(real_points, observations, rotation_matrix, K):
    """
    K: 相机内参
    database: 三维点库
    real_points: 观测到的二维点
    """
    assert len(real_points) == len(observations)
    assert observations.shape[1] == 2
    assert real_points.shape[1] == 3
    # 求解相机位姿

    def residual_function(params, ob, p):
        # 定义残差函数
        T = params
        uv = cv2.projectPoints(p, rotation_matrix, T, K, None)[0].squeeze()
        return (uv - ob).flatten()

    initial_pose = np.zeros(3)
    result = least_squares(
        residual_function, initial_pose, args=(observations, real_points)
    )
    if not result.success:
        return None
    return result.x  # 提取平移部分


method = [linear_least]


def single_test(n, f, d, reso, noise, drift, m=100, height=300, outlier=0):
    """
    f: 焦距/mm
    d: 像元尺寸/mm
    reso: 图像分辨率/pix
    仅使用陨石坑完成位置和姿态的估计解算
    """
    # 生成随机三维点库
    # np.random.seed(114514)
    x = np.random.uniform(-300, 300, 1000) * 1e3
    y = np.random.uniform(-300, 300, 1000) * 1e3
    z = np.random.uniform(-2, 2, 1000) * 1e3
    points_3d = np.vstack((x.flatten(), y.flatten(), z.flatten()))
    # 相机内参
    K = np.array([[f / d * 1e3, 0, reso / 2], [0, f / d * 1e3, reso / 2], [0, 0, 1]])
    err = []
    ind_sum = []
    test_num = 0
    progbar = tqdm.tqdm(total=n)
    failed = 0
    while test_num < n:
        # 随机选定旋转角度生成旋转矩阵
        Rvec = np.random.randn(3) * np.pi / 5
        pose = np.array([*np.random.uniform(-8, 8, 2), height]) * 1e3
        # ori = np.array(line[3:])
        # pose = np.array(line[:3])
        R = cv2.Rodrigues(Rvec)[0]
        # 随机生成平移向量
        T = -R @ pose
        # 生成投影矩阵
        # P = np.hstack((R, T[:, None]))
        # S = np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1]])
        # 生成二维点
        observations = (
            cv2.projectPoints(points_3d.T, R, T, K, np.zeros(5))[0].squeeze().T
        )
        # 框出图像
        # observations = points_2d[:2] / points_2d[2]
        ind = (
            (observations[0] >= 0)
            & (observations[0] < reso)
            & (observations[1] >= 0)
            & (observations[1] < reso)
        )
        if ind.sum() != m:
            continue
        observations = observations[:, ind].T
        real_points = points_3d[:, ind].T
        # 加入噪声和漂移
        observations += np.random.randn(*observations.shape) * noise + drift
        # 加入外点
        ind = np.random.rand(m) < outlier
        observations[ind] = np.random.uniform(0, reso, (ind.sum(), 2))
        # 估计位姿
        P_ba = initial_pose(K, observations, real_points)
        if P_ba[0] is None:
            failed += 1
            test_num += 1
            progbar.update(1)
            continue
        # P_ba = bundle_adjustment(K @ S, observations, real_points, P_ba)
        # 提取位置和姿态
        ## TODO 消除二义性！！这一步巨坑
        position = P_ba[1]  # 提取平移部分
        rotation_matrix = P_ba[0]  # 提取旋转部分
        esti_post = -(rotation_matrix.T @ position).squeeze()
        if esti_post[-1] < 0:
            rotation_matrix[:, 2] = -rotation_matrix[:, 2]
            esti_post = -(rotation_matrix.T @ position).squeeze()
        #     rot = (S @ rotation_matrix).T
        else:
            rotation_matrix = rotation_matrix / np.linalg.det(rotation_matrix)
        # 估计误差
        err_dist = esti_post - pose
        err_rot = cv2.Rodrigues(rotation_matrix.T @ R)[0] * 180 / np.pi
        err.append((err_dist, err_rot))
        test_num += 1
        progbar.update(1)
    err = np.array(err)
    pose_err = np.sqrt(np.power(err[:, 0], 2).mean(axis=0))
    rot_err = np.sqrt(np.power(err[:, 1], 2).mean(axis=0))
    with open(f"ind_sum_{f}.csv", "w") as w:
        for p, r, i in zip(err[:, 0], err[:, 1], ind_sum):
            print(",".join(map(str, chain(p, r, i))), file=w, flush=True)
    print(f"Error Type\tx\ty\tz\ttotal\tnoise:{noise}\tdrift:{drift}\toutlier:{outlier}")
    print(
        "Position(m)\t"
        + "\t".join(f"{err:.3f}" for err in pose_err)
        + f"\t{np.linalg.norm(pose_err):.3f}"
    )
    print(
        "Rotation(deg)\t"
        + "\t".join(f"{err:.3f}" for err in rot_err)
        + f"\t{np.linalg.norm(rot_err):.3f}"
    )
    print(f"Success Rate: {1-failed/n:.3f}")


def simulate_points(height, num_points):
    """
    模拟在高度为h km的近低轨道上，所能看到的月球表面点云
    Arguments:
      height: 高度/km
      num_points: 生成点云数量
    """
    # 求解纬度角
    radius = 1737.4  # km
    theta = height / radius
    theta = np.arctan(theta * theta + 2 * theta)
    # 生成点云
    np.random.seed(114514)
    theta = np.random.uniform(0, theta, num_points)
    phi = np.random.uniform(0, 2 * np.pi, num_points)
    dr = np.random.uniform(-2, 2, num_points)
    x = (radius + dr) * np.sin(theta) * np.cos(phi)
    y = (radius + dr) * np.sin(theta) * np.sin(phi)
    z = (radius + dr) * np.cos(theta) - radius
    return np.vstack((x, y, z)) * 1e3


def single_test_combine_ori(
    n, f, d, reso_x, reso_y, noise, drift, outlier, m=50, height=300
):
    """
    f: 焦距/mm
    d: 像元尺寸/mm
    reso_x: 图像x轴分辨率/pix
    reso_y: 图像y轴分辨率/pix
    在利用星敏感器已知姿态的情况下，使用陨石坑完成位置和姿态的估计解算
    """
    # 生成随机三维点库
    # np.random.seed(114514)
    # x = np.random.uniform(-300, 300, 1000) * 1e3
    # y = np.random.uniform(-300, 300, 1000) * 1e3
    # z = np.random.uniform(-2, 2, 1000) * 1e3
    # points_3d = np.vstack((x.flatten(), y.flatten(), z.flatten()))

    # 1066是由积分估计出来的
    points_3d = simulate_points(height, 1066)
    # 相机内参
    K = np.array(
        [[f / d * 1e3, 0, reso_x / 2], [0, f / d * 1e3, reso_y / 2], [0, 0, 1]]
    )
    err = {me.__name__: [] for me in method}
    test_num = 0
    # m = potential_points(height)
    failed = {me.__name__: 0 for me in method}
    progbar = tqdm.tqdm(total=n)
    while test_num < n:
        # line = list(map(float, line))
        # 随机选定旋转角度生成旋转矩阵
        ori = np.random.randn(3) * np.pi / 20
        pose = np.array([0, 0, height]) * 1e3
        # ori = np.array(line[3:])
        # pose = np.array(line[:3])
        R = Rotation.from_rotvec(ori).as_matrix().T
        # 随机生成平移向量
        T = -R @ pose
        # 生成投影矩阵
        P = np.hstack((R, T[:, None]))
        S = np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1]])
        # 生成二维点
        points_2d = (
            K
            @ S
            @ P
            @ np.pad(points_3d, ((0, 1), (0, 0)), "constant", constant_values=1)
        )
        # 框出图像
        observations = points_2d[:2] / points_2d[2]
        ind = (
            (observations[0] >= 0)
            & (observations[0] < reso_x)
            & (observations[1] >= 0)
            & (observations[1] < reso_y)
        )
        if ind.sum() <= m:
            failed = {k: v + 1 for k, v in failed.items()}
            progbar.update(1)
            test_num += 1
            continue
        observations = observations[:, ind].T
        real_points = points_3d[:, ind].T
        # 加入噪声
        observations += np.random.randn(*observations.shape) * noise + drift
        # 加入外点
        ind = np.random.rand(ind.sum()) < outlier
        observations[ind] = np.random.uniform(0, reso_y, (ind.sum(), 2))
        # 星敏感器输出的姿态精度为60角秒，相当于在角度上增加此噪声
        esti_ori = ori + np.random.randn(3) * np.pi / 180 / 3600 * 16
        rotation_matrix = Rotation.from_rotvec(esti_ori).as_matrix().T
        # 最小二乘法直接可解出平移向量, K @ R @ X_i + K @ T = s * x_i
        # 解得结果为 [inv(K) @ x_i]x @ T = [inv(K) @ x_i]x @ K @ R @ X_i
        for me in method:
            # 提取位置和姿态
            translation = me(real_points, observations, rotation_matrix, K)
            if translation is None:
                failed[me.__name__] += 1
                continue
            ## TODO 消除二义性！！这一步巨坑
            esti_post = (rotation_matrix.T @ translation).squeeze()
            if esti_post[-1] > 0:
                rotation_matrix[:, 2] = -rotation_matrix[:, 2]
                esti_post = (rotation_matrix.T @ translation).squeeze()
            # 估计误差
            err_dist = esti_post + pose
            err_rot = (
                Rotation.from_matrix(rotation_matrix) * Rotation.from_rotvec(ori)
            ).as_rotvec(True)
            err[f"{me.__name__}"].append((err_dist, err_rot))
        progbar.update(1)
        test_num += 1
    progbar.close()
    line = ""
    for me in method:
        err_ = np.array(err[f"{me.__name__}"])
        pose_err = np.sqrt(np.power(err_[:, 0], 2).mean(axis=0))
        rot_err = np.sqrt(np.power(err_[:, 1], 2).mean(axis=0))
        print("Error Type\tx\ty\tz\ttotal\tfor method:", me.__name__,f"\tnoise:{noise}\tdrift:{drift}\toutlier:{outlier}")
        print(
            "Position(m)\t"
            + "\t".join(f"{e:.3f}" for e in pose_err)
            + f"\t{np.linalg.norm(pose_err):.3f}"
        )
        print(
            "Rotation(deg)\t"
            + "\t".join(f"{e:.3f}" for e in rot_err)
            + f"\t{np.linalg.norm(rot_err):.3f}"
        )
        print(f"Success Rate: {1-failed[me.__name__]/n:.3f}")
        line += ",".join(map(str, rot_err)) + "," + ",".join(map(str, pose_err))

    return line


if __name__ == "__main__":
    args = arg_parse()
    with open(f"result-{args.f}-{args.d}-{args.reso_x}.csv", "a") as w:
        line = single_test_combine_ori(**args.__dict__)
        angle_x = np.arctan(args.d / args.f / 1e3 / 2 * args.reso_x) * 360 / np.pi
        angle_y = np.arctan(args.d / args.f / 1e3 / 2 * args.reso_y) * 360 / np.pi
        w.write(
            f"{args.height},{args.d},{args.reso_x},{args.reso_y},{args.f},{angle_x},{angle_y},{args.noise},{args.drift},{line},{args.m}\n"
        )
